import unittest
from app.dao import RagDAO

class TestRagDAO(unittest.TestCase):
    def setUp(self):
        """设置每个测试方法之前的准备工作"""
        # Mock `parse_args` to avoid system exit during tests

        self.rag_dao = RagDAO()
        self.title = "BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding"
        self.abstract = (
            "We introduce a new language representation model called BERT, which stands for Bidirectional Encoder "
            "Representations from Transformers. Unlike recent language representation models, BERT is designed to pre-train "
            "deep bidirectional representations from unlabeled text by jointly conditioning on both left and right context in "
            "all layers. As a result, the pre-trained BERT model can be fine-tuned with just one additional output layer to "
            "create state-of-the-art models for a wide range of tasks, such as question answering and language inference, "
            "without substantial task-specific architecture modifications. BERT is conceptually simple and empirically powerful. "
            "It obtains new state-of-the-art results on eleven natural language processing tasks, including pushing the GLUE score "
            "to 80.5% (7.7% point absolute improvement), MultiNLI accuracy to 86.7% (4.6% absolute improvement), SQuAD v1.1 question "
            "answering Test F1 to 93.2 (1.5 point absolute improvement) and SQuAD v2.0 Test F1 to 83.1 (5.1 point absolute improvement)."
        )

    def test_search_milvus(self):
        """测试 Milvus 查询是否返回正确的格式"""
        results = self.rag_dao.search_milvus(self.title, self.abstract, limit=5)

        # 验证返回的结果格式正确
        self.assertGreater(len(results), 0)
        self.assertEqual(len(results[0]), 4)  # 确保每个结果是一个包含 4 个元素的元组

    def test_search_semanticscholar(self):
        """测试 Semantic Scholar 查询是否返回正确的格式"""
        # 模拟 Semantic Scholar API 返回

        results = self.rag_dao.search_semanticscholar(self.title, limit=5)

        # 验证返回的结果格式正确
        self.assertGreater(len(results), 0)
        self.assertEqual(len(results[0]), 4)  # 确保每个结果是一个包含 4 个元素的元组
        
    def test_search_crossref(self):
        """测试 Crossref 查询是否返回正确的格式"""
        # 模拟 Crossref API 返回

        results = self.rag_dao.search_crossref(self.title, limit=5)

        # 验证返回的结果格式正确
        self.assertGreater(len(results), 0)
        self.assertEqual(len(results[0]), 4)  # 确保每个结果是一个包含 4 个元素的元组


if __name__ == "__main__":
    unittest.main()
